import pandas as pd
import numpy as np
from datetime import datetime
from sklearn.preprocessing import LabelEncoder

def load_data(file_path):
    """
    加载原始犯罪数据
    """
    df = pd.read_csv(file_path)
    print(f"Loaded {len(df)} records")
    return df

def preprocess_data(df):
    """
    数据预处理：清洗数据，处理缺失值，转换数据类型
    """
    # 复制数据框以避免修改原始数据
    df = df.copy()
    
    # 转换日期时间
    df['DATE OCC'] = pd.to_datetime(df['DATE OCC'])
    
    # 提取时间特征
    df['Year'] = df['DATE OCC'].dt.year
    df['Month'] = df['DATE OCC'].dt.month
    df['Day'] = df['DATE OCC'].dt.day
    df['Hour'] = df['TIME OCC'].apply(lambda x: int(str(x).zfill(4)[:2]))
    
    # 处理缺失值
    df = df.dropna(subset=['LAT', 'LON'])  # 删除没有位置信息的记录
    
    # 编码分类特征
    le = LabelEncoder()
    categorical_columns = ['Crm Cd Desc', 'AREA NAME', 'Premis Desc', 'Weapon Desc']
    for col in categorical_columns:
        if col in df.columns:
            df[f'{col}_encoded'] = le.fit_transform(df[col].fillna('UNKNOWN'))
    
    # 创建时间段特征
    df['TimeOfDay'] = pd.cut(df['Hour'], 
                            bins=[0, 6, 12, 18, 24], 
                            labels=['Night', 'Morning', 'Afternoon', 'Evening'])
    
    return df

def create_spatial_features(df):
    """
    创建空间特征
    """
    # 计算与市中心的距离（以洛杉矶市政厅为中心点）
    city_hall_lat, city_hall_lon = 34.0522, -118.2437
    
    df['Distance_to_Center'] = np.sqrt(
        (df['LAT'] - city_hall_lat)**2 + 
        (df['LON'] - city_hall_lon)**2
    )
    
    # 创建网格化的位置特征
    df['Lat_Grid'] = pd.qcut(df['LAT'], q=20, labels=False)
    df['Lon_Grid'] = pd.qcut(df['LON'], q=20, labels=False)
    
    return df

def aggregate_crime_stats(df):
    """
    聚合统计特征
    """
    # 按网格统计犯罪数量
    grid_stats = df.groupby(['Lat_Grid', 'Lon_Grid']).size().reset_index(name='Crime_Count')
    
    # 按时间段统计
    time_stats = df.groupby(['Year', 'Month', 'Lat_Grid', 'Lon_Grid']).size().reset_index(name='Monthly_Crime_Count')
    
    return grid_stats, time_stats

def save_processed_data(df, output_path):
    """
    保存处理后的数据
    """
    df.to_csv(output_path, index=False)
    print(f"Saved processed data to {output_path}")

if __name__ == "__main__":
    # 设置输入输出路径
    input_path = "data/raw/crime-data-los-angeles.csv"
    output_path = "data/processed/processed_crime_data.csv"
    
    # 加载数据
    df = load_data(input_path)
    
    # 数据预处理
    df = preprocess_data(df)
    
    # 创建空间特征
    df = create_spatial_features(df)
    
    # 保存处理后的数据
    save_processed_data(df, output_path)
    
    print("Data processing completed successfully!") 